package com.cyw.common.config.trace;


import com.cyw.common.constants.CommonConstant;
import org.apache.commons.lang.StringUtils;
import org.apache.dubbo.common.constants.CommonConstants;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

/**
 * dubbo服务间的traceId传递
 */
@Activate(group = {CommonConstants.CONSUMER, CommonConstants.PROVIDER}, order = -9999)
public class MDCTraceDubboFilter implements Filter {

    private static final Logger LOGGER = LoggerFactory.getLogger(MDCTraceDubboFilter.class);

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        String traceIdKey = CommonConstant.traceIdKey;

        // 消费者
        if (RpcContext.getContext().isConsumerSide()) {
            String traceId = MDC.get(traceIdKey);

            // 消费者调用之前应该已经设置过traceId了，否则生成traceId（这个traceId只会出现在provider中了）
            if (StringUtils.isBlank(traceId)) {
                LOGGER.warn("Dubbo consumer side, traceId is blank! You should set it in advance!");
                traceId = TraceIdGenerator.traceId();
            }

            // 传递traceId
            RpcContext.getContext().setAttachment(traceIdKey, traceId);
        }

        // 提供者
        if (RpcContext.getContext().isProviderSide()) {
            String traceId = RpcContext.getContext().getAttachment(traceIdKey);

            // 消费者未传递traceId
            if (StringUtils.isBlank(traceId)) {
                LOGGER.warn("Provider side, traceId is blank! You should check counsumer!");
                traceId = TraceIdGenerator.traceId();
            }
            MDC.put(traceIdKey, traceId);
        }
        return invoker.invoke(invocation);
    }
}
